Summary

Fixes #7037.

Dataset.to_json(lines=False) raised NotImplementedError (and before the guard was added, silently produced invalid JSON) whenever batch_size < num_rows — i.e., the default case for any dataset larger than DEFAULT_MAX_BATCH_SIZE (1000 rows).

Root cause: Each batch was serialised to its own JSON array ([{...}, ...]) and written back-to-back in the file, producing e.g. [{...}][{...}] — syntactically invalid.

Fix (two changes in src/datasets/io/json.py):

1. write() — Narrow the NotImplementedError guard to dict-producing orientations (split, index, columns, table) only. Array-producing orientations (records, values) can now be batched safely.

2. _write() — When lines=False and orient in ("records", "values"), strip the outer [ / ] from each batch's bytes and emit a single wrapping [...], streaming with no extra memory overhead. Dict-producing orientations fall through to the unchanged write-as-is path.

As a bonus, the single-proc and multiprocessing iteration loops are unified via contextlib.nullcontext, removing ~10 lines of duplication.

Verified:

import io, json
from datasets import Dataset
from datasets.io.json import JsonDatasetWriter

ds = Dataset.from_dict({"id": list(range(25)), "value": [f"v{i}" for i in range(25)]})
buf = io.BytesIO()
JsonDatasetWriter(ds, buf, batch_size=10, lines=False).write()  # 3 batches
buf.seek(0)
result = json.loads(buf.read())   # previously raised NotImplementedError / produced invalid JSON
assert len(result) == 25          # single valid JSON array with all rows

Test plan

• New parametrised test test_dataset_to_json_lines_false_batched covers orient="records" and orient="values" with batch_size < num_rows
• Full TestJsonDatasetWriter suite passes (20 tests; 3 pre-existing errors for shared_datadir fixture from missing pytest-datadir plugin, unrelated to this change)